package com.shanzhi.connection.common;

import com.shanzhi.connection.common.entity.RSAKeyEntity;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import java.io.*;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

public class RSASecurityUtils {


    /**
     * 加密
     *
     * @param key KEY
     * @param in  输入参数
     * @param out 输出加密后的密文
     * @throws NoSuchAlgorithmException
     * @throws NoSuchPaddingException
     * @throws InvalidKeyException
     * @throws IOException
     * @throws BadPaddingException
     * @throws IllegalBlockSizeException
     */
    public static void encode(Key key, InputStream in, OutputStream out) throws NoSuchAlgorithmException,
            NoSuchPaddingException, InvalidKeyException, IOException, IllegalBlockSizeException, BadPaddingException {
        // 最大的加密明文长度
        final int maxEncryptBlock = 245;

        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        Cipher cipher = Cipher.getInstance(keyFactory.getAlgorithm());
        cipher.init(Cipher.ENCRYPT_MODE, key);

        byte[] buffer = new byte[maxEncryptBlock];
        int len = 0;
        while ((len = in.read(buffer)) != -1) {
            out.write(cipher.doFinal(buffer, 0, len));
        }
    }

    /**
     * 解密
     *
     * @param key KEY
     * @param in  输入参数
     * @param out 输出解密后的原文
     * @throws NoSuchAlgorithmException
     * @throws NoSuchPaddingException
     * @throws InvalidKeyException
     * @throws IOException
     * @throws BadPaddingException
     * @throws IllegalBlockSizeException
     */
    public static void decode(Key key, InputStream in, OutputStream out) throws NoSuchAlgorithmException,
            NoSuchPaddingException, InvalidKeyException, IOException, IllegalBlockSizeException, BadPaddingException {

        // 最大的加密明文长度
        final int maxDecryptBlock = 256;

        KeyFactory keyFactory = KeyFactory.getInstance("RSA");
        Cipher cipher = Cipher.getInstance(keyFactory.getAlgorithm());
        cipher.init(Cipher.DECRYPT_MODE, key);

        byte[] buffer = new byte[maxDecryptBlock];
        int len = 0;
        while ((len = in.read(buffer)) != -1) {
            out.write(cipher.doFinal(buffer, 0, len));
        }
    }


    /**
     * 从 RSAKeyEntity 获取 RSAPublicKey 对象
     *
     * @param rsaKeyEntity 包含 Base64 编码的密钥信息
     * @return RSAPublicKey
     * @throws Exception
     */
    public static RSAPublicKey getPublicKeyFromKeyStore(RSAKeyEntity rsaKeyEntity) throws Exception {
        String base64PublicKey = rsaKeyEntity.getPublicKey();
        byte[] publicKeyBytes = Base64.getDecoder().decode(base64PublicKey);

        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(publicKeyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");

        return (RSAPublicKey) keyFactory.generatePublic(keySpec);
    }

    /**
     * 从 RSAKeyEntity 获取 RSAPrivateKey 对象
     *
     * @param rsaKeyEntity 包含 Base64 编码的密钥信息
     * @return RSAPrivateKey
     * @throws Exception
     */
    public static RSAPrivateKey getPrivateKeyFromKeyStore(RSAKeyEntity rsaKeyEntity) throws Exception {
        String base64PrivateKey = rsaKeyEntity.getPrivateKey();
        byte[] privateKeyBytes = Base64.getDecoder().decode(base64PrivateKey);

        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(privateKeyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance("RSA");

        return (RSAPrivateKey) keyFactory.generatePrivate(keySpec);
    }

    public static RSAKeyEntity buildBase64Keys() throws NoSuchAlgorithmException {

        // 生成 RSA 密钥对
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
        keyPairGenerator.initialize(2048);
        KeyPair keyPair = keyPairGenerator.generateKeyPair();


        // 公钥和私钥
        RSAPublicKey rsaPublicKey = (RSAPublicKey) keyPair.getPublic();
        RSAPrivateKey rsaPrivateKey = (RSAPrivateKey) keyPair.getPrivate();

        return new RSAKeyEntity(Base64.getEncoder().encodeToString(rsaPrivateKey.getEncoded()), Base64.getEncoder().encodeToString(rsaPublicKey.getEncoded()));
    }


    /**
     * 私钥  ---   加密字符串    --   用于需要看见的明文密钥
     *
     * @return
     */
    public static String encodeStr(String originStr, RSAPublicKey rsaPublicKey) {
        // 要加密的原文
        byte[] content = originStr.getBytes(StandardCharsets.UTF_8);
        System.out.println("原文：" + new String(content, StandardCharsets.UTF_8));

        // 加密后的密文
        ByteArrayOutputStream encryptedout = new ByteArrayOutputStream();
        // 公钥加密
        try {
            encode(rsaPublicKey, new ByteArrayInputStream(content), encryptedout);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }

        System.out.println("加密后的密文：" + Base64.getEncoder().encodeToString(encryptedout.toByteArray()));
        return Base64.getEncoder().encodeToString(encryptedout.toByteArray());
    }


    /**
     * 公钥  ---   解密字符串    --   用于需要看不见的密钥
     * @param encodeStr
     * @param rsaPrivateKey
     * @return
     * @throws NoSuchPaddingException
     * @throws IllegalBlockSizeException
     * @throws NoSuchAlgorithmException
     * @throws IOException
     * @throws BadPaddingException
     * @throws InvalidKeyException
     */
    public static String decodeStr(String encodeStr, RSAPrivateKey rsaPrivateKey)  {
        // 解密后的原文
        ByteArrayOutputStream decryptedOut = new ByteArrayOutputStream();

        // 使用私钥解密
        try {
            decode(rsaPrivateKey, new ByteArrayInputStream(Base64.getDecoder().decode(encodeStr)), decryptedOut);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }

        return new String(decryptedOut.toByteArray(), StandardCharsets.UTF_8);
    }

    public static void main(String[] args) throws Exception {
        RSAKeyEntity rsaKeyEntity = buildBase64Keys();
        RSAPublicKey rsaPublicKey = getPublicKeyFromKeyStore(rsaKeyEntity);
        RSAPrivateKey rsaPrivateKey = getPrivateKeyFromKeyStore(rsaKeyEntity);

        // 加密的原文
        String encodeStr = encodeStr("szsz", rsaPublicKey);

        // 解密后的原文
        String s = decodeStr(encodeStr, rsaPrivateKey);
    }
}
